from package.emotionDetector_interface import EmotionDetector
# Modal Loader
from fastai.vision.all import *  
import numpy  

class EmotionsDeepModal(EmotionDetector):
    """
    Implements the detector of prototypic emotions on face images using the
    DeepModal library.
    """
    trained_modal_labels = {'Angry': 0, 'Happy': 1, 'Neutral': 2, 'Sad': 3, 'Surprise': 4, 'Ahegao': 5}  
    # 创建一个逆映射，将索引映射到情感标签  
    inverse_labels = {v: k for k, v in trained_modal_labels.items()}  
    #---------------------------------------------
    def __init__(self):
        """
        Class constructor.
        """
        self.modal : 'fastai.learner.Learner' = None
        self.load()
        super().__init__()
    def load(self) -> bool:  
        # Get the current working directory  
        cwd = os.getcwd()  
        modalFilePath = os.path.join(cwd, 'models', 'deep_model.pkl') 
        try: 
            self.modal = load_learner(modalFilePath)  
        except Exception as e:  # It's better to catch a specific exception, or at least print the error  
            print("Load Model Failed:", e)  
            return False  # Return False if loading fails  
        
        # No need for an if-else here, just return True since we know the model loaded successfully  
        return self.modal is not None
    def _get_emotion_probabilities(self, outputs: 'torch.Tensor') -> 'OrderedDict':  
        # 创建一个OrderedDict来保存结果  
        emotion_probabilities = OrderedDict()  
        
        # 遍历outputs，将其与情感标签关联，并添加到OrderedDict中  
        for i, prob in enumerate(outputs):  
            emotion_label = self.inverse_labels.get(i, 'Unknown')  # 如果索引不在labels中，则默认为'Unknown'  
            # emotion_label = emotion_label.lower()  # 将情感标签转换为小写以匹配期望的输出格式  
            emotion_probabilities[emotion_label] = prob.item()  
        
        return emotion_probabilities  
    def detect(self,image: 'numpy.array') -> 'OrderedDict':
        try: 
            pred_class, pred_idx, outputs =  self.modal.predict(image)
            return self._get_emotion_probabilities(outputs)  
            # return np.random(1,2,1)
        except Exception as e:
            print("Predict Model Failed:", e)  
            return False  

if __name__ == "__main__":  
    instance  = EmotionsDeepModal()
    help(instance)
    